#include "identtrackerorb.h"
#include <opencv2/calib3d.hpp>
#include "source/logger.h"

IdentTrackerORB::IdentTrackerORB(float scaleFactor) : IdentTracker(scaleFactor), minDistance(DBL_MAX), similarityThreshold(0.75)
{
    orb = cv::ORB::create();
}

void IdentTrackerORB::initialize(const cv::Mat &templateImg, const cv::Rect& roi)
{
    LogDebug << "Initializing IdentTrackerORB";
    computeTemplateDescriptors(templateImg);
    bestIdentifiedRect = roi; // Reset best identified rect
    minDistance = DBL_MAX; // Reset min distance
}


cv::Rect2d IdentTrackerORB::detectAndMatch(const cv::Mat &frame, float &Similarity)
{
    std::vector<cv::KeyPoint> keypointsFrame;
    cv::Mat descriptorsFrame;
    orb->detectAndCompute(frame, cv::Mat(), keypointsFrame, descriptorsFrame);

    cv::BFMatcher matcher(cv::NORM_HAMMING);
    std::vector<cv::DMatch> matches;

    // Ratio test for filtering matches
    std::vector<std::vector<cv::DMatch>> knn_matches;
    matcher.knnMatch(descriptorsTemplate, descriptorsFrame, knn_matches, 2);
    for (const auto &knn_match : knn_matches)
    {
        if (knn_match[0].distance < similarityThreshold * knn_match[1].distance)
        {
            matches.push_back(knn_match[0]);
        }
    }

    if (matches.empty())
    {
        LogWarning << "No matches found in IdentTrackerORB";
        return cv::Rect2d();
    }

    std::vector<cv::Point2f> pointsTemplate, pointsFrame;
    for (const auto &match : matches)
    {
        pointsTemplate.push_back(keypointsTemplate[match.queryIdx].pt);
        pointsFrame.push_back(keypointsFrame[match.trainIdx].pt);
    }

    cv::Mat homography;
    try
    {
        homography = cv::findHomography(pointsTemplate, pointsFrame, cv::RANSAC);
    }
    catch (const cv::Exception &e)
    {
        LogError << "cv::findHomography failed: " << e.what();
        return cv::Rect2d();
    }

    if (homography.empty() || cv::countNonZero(homography) < 3)
    {
        LogWarning << "Homography calculation failed in IdentTrackerORB";
        return cv::Rect2d();
    }

    std::vector<cv::Point2f> identPoints = {
                                            {0.0f, 0.0f},
                                            {static_cast<float>(descriptorsTemplate.cols), 0.0f},
                                            {static_cast<float>(descriptorsTemplate.cols), static_cast<float>(descriptorsTemplate.rows)},
                                            {0.0f, static_cast<float>(descriptorsTemplate.rows)}};
    std::vector<cv::Point2f> transformedPoints;
    cv::perspectiveTransform(identPoints, transformedPoints, homography);

    // Debug information
    for (size_t i = 0; i < identPoints.size(); ++i)
    {
        LogDebug << "IdentPoint " << i << ": (" << identPoints[i].x << ", " << identPoints[i].y << ")";
        LogDebug << "TransformedPoint " << i << ": (" << transformedPoints[i].x << ", " << transformedPoints[i].y << ")";
    }

    // Check if the transformed points form a valid rectangle
    double width = cv::norm(transformedPoints[1] - transformedPoints[0]);
    double height = cv::norm(transformedPoints[2] - transformedPoints[1]);
    LogDebug << "Transformed width: " << width << ", height: " << height;

    if (width < 10 || height < 10)
    {
        LogWarning << "Invalid transformed points, width or height too small";
        return cv::Rect2d();
    }

    cv::Rect2d identifiedRect = cv::boundingRect(transformedPoints);

    // Calculate the distance to the previous identified location
    cv::Point2f prevCenter(prevRect.x + prevRect.width / 2, prevRect.y + prevRect.height / 2);
    cv::Point2f currCenter(identifiedRect.x + identifiedRect.width / 2, identifiedRect.y + identifiedRect.height / 2);
    double distance = cv::norm(prevCenter - currCenter);

    // Select the closest identification
    if (distance < minDistance)
    {
        minDistance = distance;
        bestIdentifiedRect = identifiedRect;
    }

    return bestIdentifiedRect;
}

void IdentTrackerORB::computeTemplateDescriptors(const cv::Mat &templateImg)
{
    orb->detectAndCompute(templateImg, cv::Mat(), keypointsTemplate, descriptorsTemplate);
    LogDebug << "Template descriptors computed.";
}
